package com.flute.haflute.jobbox.base.storage;

import static me.prettyprint.cassandra.utils.StringUtils.bytes;
import static me.prettyprint.cassandra.utils.StringUtils.string;

import java.io.UnsupportedEncodingException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.WeakHashMap;

import me.prettyprint.cassandra.service.CassandraClientPool;
import me.prettyprint.cassandra.service.Keyspace;

import org.apache.cassandra.thrift.Column;
import org.apache.cassandra.thrift.ColumnParent;
import org.apache.cassandra.thrift.ColumnPath;
import org.apache.cassandra.thrift.KeyRange;
import org.apache.cassandra.thrift.NotFoundException;
import org.apache.cassandra.thrift.SlicePredicate;
import org.apache.cassandra.thrift.SliceRange;
import org.apache.cassandra.thrift.UnavailableException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.flute.haflute.jobbox.base.exception.DataNotFoundException;
import com.flute.haflute.jobbox.base.exception.ServiceUnavailableException;

public class CassandraStorageServiceImpl implements KeyValueStorageService {
	private static Logger logger = LoggerFactory
			.getLogger(CassandraStorageServiceImpl.class);

	private final String COLUMN = "content";

	private boolean localCassandra = false;

	private boolean borrowed = false;

	private Keyspace keyspace;

	private ColumnPath columnPath;

	private CassandraRangeQuery rangeQuerier;

	public CassandraStorageServiceImpl(Keyspace keyspace,
			ColumnPath columnPath, boolean localCassandra) {
		this.keyspace = keyspace;
		this.columnPath = columnPath;
		this.localCassandra = localCassandra;
		columnPath.setColumn(bytes(COLUMN));
		rangeQuerier = new CassandraRangeQuery(keyspace, new ColumnParent(
				columnPath.column_family));
	}

	/**
	 * @return the columnFamily
	 */
	public String getColumnFamily() {
		return columnPath.column_family;
	}

	@Override
	public synchronized void insert(String key, String value)
			throws ServiceUnavailableException {
		try {
			keyspace.insert(key, columnPath, bytes(value));
		} catch (Exception e) {
			logger.error("insert one colum record Key=" + key, e);
			ServiceUnavailableException sue = new ServiceUnavailableException(
					e.getStackTrace());
			throw sue;
		}

	}

	@Override
	public synchronized void insert(byte[] key, byte[] value)
			throws ServiceUnavailableException {
		try {
			keyspace.insert(new String(key, "utf-8"), columnPath, value);
		} catch (Exception e) {
			logger.error("insert one colum record Key=" + string(key), e);
			ServiceUnavailableException sue = new ServiceUnavailableException(
					e.getStackTrace());
			throw sue;
		}
	}

	/*
	 * @see
	 * com.flute.jobbox.base.storage.KeyValueStorageService#batchInsert(java
	 * .util.Map)
	 */
	@Override
	public synchronized void batchInsert(Map<byte[], byte[]> keyvalueMap)
			throws ServiceUnavailableException {
		Map<String, List<Column>> cfmap = new HashMap<String, List<Column>>();
		for (Map.Entry<byte[], byte[]> entry : keyvalueMap.entrySet()) {
			ArrayList<Column> columnList = new ArrayList<Column>(1);
			Column col = new Column();
			col.setName(bytes("content"));
			col.setValue(entry.getValue());
			columnList.add(col);
			try {
				cfmap.put(new String(entry.getKey(), "utf-8"), columnList);
			} catch (UnsupportedEncodingException e) {
				e.printStackTrace();
			}
		}
		try {
			keyspace.batchInsert("haflute7", cfmap, null);
		} catch (Exception e) {
			logger.error("insert batch columns record Key", e);
			ServiceUnavailableException sue = new ServiceUnavailableException(
					e.getStackTrace());
			throw sue;
		}
	}

	/*
	 * @see com.flute.jobbox.base.storage.StorageService#get(java.lang.String)
	 */
	@Override
	public String get(String key) throws DataNotFoundException,
			ServiceUnavailableException {
		try {
			Column col = keyspace.getColumn(key, columnPath);
			return string(col.getValue());
		} catch (NotFoundException ne) {
			logger.error("get one colum record Key=" + key, ne);
			DataNotFoundException dne = new DataNotFoundException(
					ne.getStackTrace());
			throw dne;
		} catch (Exception e) {
			logger.error("get one colum record Key=" + key, e);
			ServiceUnavailableException sue = new ServiceUnavailableException(
					e.getStackTrace());
			throw sue;
		}
	}

	/*
	 * @see com.flute.jobbox.base.storage.StorageService#get(byte[])
	 */
	@Override
	public byte[] get(byte[] key) throws ServiceUnavailableException,
			DataNotFoundException {
		try {
			Column col = keyspace.getColumn(new String(key), columnPath);
			return col.getValue();
		} catch (NotFoundException ne) {
			logger.error("get one colum record Key=" + string(key), ne);
			DataNotFoundException dne = new DataNotFoundException(
					ne.getStackTrace());
			throw dne;
		} catch (Exception e) {
			logger.error("get one colum record Key=" + string(key), e);
			ServiceUnavailableException sue = new ServiceUnavailableException(
					e.getStackTrace());
			throw sue;
		}
	}

	protected boolean isLocalCassandra() {
		return localCassandra;
	}

	protected void markAsBorrowed() {
		borrowed = true;
	}

	protected void markAsUnborrowed() {
		borrowed = false;
	}

	protected boolean canBorrow() {
		return !borrowed;
	}

	protected void releaseClient(CassandraClientPool pool) {
		try {
			pool.releaseClient(keyspace.getClient());
			pool.releaseKeyspace(keyspace);
		} catch (Exception e) {
			logger.info("release Client error", e);
		}
	}

	@Override
	public Map<String, String> batchGet(String begin, String end)
			throws ServiceUnavailableException {
		try {
			return rangeQuerier.queryColumn(begin, end, COLUMN);
		} catch (Exception e) {
			logger.error("", e);
			throw new ServiceUnavailableException(e.getStackTrace());
		}
	}

	@Override
	public String serviceType() {
		return getColumnFamily();
	}

	@Override
	public void batchDelete(List<String> keys)
			throws ServiceUnavailableException {
		for (String key : keys) {
			try {
				keyspace.remove(key, columnPath);
			} catch (UnavailableException e) {
				throw new ServiceUnavailableException(e.getStackTrace());
			} catch (Exception e) {
				logger.error("batch delete error, see details from Cassandra,",
						e);
			}
		}
	}

	@Override
	public void delete(String key) throws DataNotFoundException,
			ServiceUnavailableException {
		try {
			keyspace.remove(key, columnPath);
		} catch (Exception e) {
			throw new ServiceUnavailableException(e.getStackTrace());
		}
	}

	@Override
	public void delete(byte[] key) throws DataNotFoundException,
			ServiceUnavailableException {
		try {
			keyspace.remove(string(key), columnPath);
		} catch (Exception e) {
			throw new ServiceUnavailableException(e.getStackTrace());
		}
	}
}

class CassandraRangeQuery {
	private static final int COUNT = 100;
	private SlicePredicate sp;
	private KeyRange keyRange;
	private Keyspace keyspace;
	private ColumnParent columnParent;

	public CassandraRangeQuery(Keyspace keyspace, ColumnParent columnParent) {
		keyRange = new KeyRange();
		keyRange.setCount(COUNT);

		sp = new SlicePredicate();
		SliceRange sr = new SliceRange();
		sr.setStart(new byte[0]);
		sr.setFinish(new byte[0]);
		sp.setSlice_range(sr);

		this.keyspace = keyspace;
		this.columnParent = columnParent;
	}

	public Map<String, Map<String, String>> query(String startKey, String endKey)
			throws Exception {
		keyRange.setStart_key(startKey.getBytes());
		keyRange.setEnd_key(endKey.getBytes());

		Map<String, List<Column>> rangeResult = keyspace.getRangeSlices(
				columnParent, sp, keyRange);

		Map<String, Map<String, String>> result = new HashMap<String, Map<String, String>>(
				rangeResult.size());
		for (Entry<String, List<Column>> entry : rangeResult.entrySet()) {
			String key = entry.getKey();
			Map<String, String> columns = new HashMap<String, String>(entry
					.getValue().size());
			for (Column col : entry.getValue()) {
				columns.put(string(col.getName()), string(col.getValue()));
			}
			result.put(key, columns);
		}

		return result;
	}

	public Map<String, String> queryColumn(String startKey, String endKey,
			String column) throws Exception {
		Map<String, Map<String, String>> rangeResult = query(startKey, endKey);
		Map<String, String> result = new WeakHashMap<String, String>(
				rangeResult.size());
		for (Entry<String, Map<String, String>> entry : rangeResult.entrySet()) {
			String key = entry.getKey();
			String value = entry.getValue().get(column);
			result.put(key, value);
		}
		return result;
	}

}
